Merge pull request #79 from Zinoex/fm/implicit_sink_state

Implicit sink state

Project.toml

@@ -1,7 +1,7 @@
 name = "IntervalMDP"
 uuid = "051c988a-e73c-45a4-90ec-875cac0402c7"
 authors = ["Frederik Baymler Mathiesen <[email protected]> and contributors"]
-version = "0.4.1"
+version = "0.4.2"
 
 [deps]
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"

ext/cuda/bellman/dense.jl

@@ -32,11 +32,11 @@ function IntervalMDP.bellman!(
     # - squeeze as many states as possible in a block
     # - use shared memory to store the values and permutation
     # - use bitonic sort to sort the values for all states in a block
-    wanted_threads = min(1024, 32 * length(Vres))
+    wanted_threads = min(1024, 32 * num_source(prob))
 
     threads = min(max_threads, wanted_threads)
     warps = div(threads, 32)
-    blocks = min(2^16 - 1, cld(length(Vres), warps))
+    blocks = min(2^16 - 1, cld(num_source(prob), warps))
     shmem =
         length(V) * (sizeof(Int32) + sizeof(Tv)) +
         warps * workspace.max_actions * sizeof(Tv)
@@ -132,7 +132,7 @@ end
     wid = fld1(threadIdx().x, warpsize())
 
     j = wid + (blockIdx().x - one(Int32)) * warps
-    @inbounds while j <= length(Vres)
+    @inbounds while j <= num_source(prob)
         state_dense_omaximization!(
             action_workspace,
             strategy_cache,

ext/cuda/bellman/sparse.jl

@@ -117,7 +117,7 @@ function try_small_sparse_bellman!(
     # TODO: Dynamically maximize the number of states per block
 
     threads = desired_warps * 32
-    blocks = min(2^16 - 1, cld(length(Vres), desired_warps))
+    blocks = min(2^16 - 1, cld(num_source(prob), desired_warps))
 
     kernel(
         workspace,
@@ -169,7 +169,7 @@ function small_sparse_bellman_kernel!(
 
     # Grid-stride loop
     j = wid + (blockIdx().x - one(Int32)) * nwarps
-    while j <= length(Vres)
+    while j <= num_source(prob)
         state_small_sparse_omaximization!(
             action_workspace,
             value_ws,
@@ -463,7 +463,7 @@ function try_large_sparse_bellman!(
     wanted_threads = min(1024, nextwarp(device(), cld(workspace.max_nonzeros, 2)))
 
     threads = min(max_threads, wanted_threads)
-    blocks = min(2^16 - 1, length(Vres))
+    blocks = min(2^16 - 1, num_source(prob))
 
     kernel(
         T1,
@@ -507,7 +507,7 @@ function large_sparse_bellman_kernel!(
 
     # Grid-stride loop
     j = blockIdx().x
-    @inbounds while j <= length(Vres)
+    @inbounds while j <= num_source(prob)
         state_sparse_omaximization!(
             action_workspace,
             value_ws,

src/models/IntervalMarkovDecisionProcess.jl

@@ -151,15 +151,23 @@ function checksize_imdp!(p::IntervalProbabilities, stateptr::AbstractVector{Int3
         throw(ArgumentError("The number of actions per state must be positive."))
     end
 
-    if num_target(p) != num_states
+    if num_states > num_target(p)
         throw(
             DimensionMismatch(
-                "The number of target states ($(num_target(p))) is not equal to the number of states in the stateptr $(num_states).",
+                "The number of target states ($(num_target(p))) is less than the number of states in the stateptr $(num_states).",
             ),
         )
     end
 
-    return Int32(num_states)
+    if stateptr[end] - 1 != num_source(p)
+        throw(
+            DimensionMismatch(
+                "The number of source states ($(num_source(p))) must be equal to the number of states in the stateptr $(num_states).",
+            ),
+        )
+    end
+
+    return Int32(num_target(p))
 end
 
 """

src/models/OrthogonalIntervalMarkovDecisionProcess.jl

@@ -116,15 +116,17 @@ function checksize_imdp!(
         throw(ArgumentError("The number of actions per state must be positive."))
     end
 
-    if prod(num_target, p.probs) != num_states
+    if num_states > prod(num_target, p.probs)
         throw(
             DimensionMismatch(
-                "The number of target states ($(prod(num_target, p.probs)) = $(map(num_target, p.probs))) is not equal to the number of states in the problem $(num_states).",
+                "The number of target states ($(prod(num_target, p.probs)) = $(map(num_target, p.probs))) is less than the number of states in the problem $(num_states).",
             ),
         )
     end
 
-    return Int32(num_states)
+    # TODO:: Check that source_dims match stateptr
+
+    return Int32(prod(num_target, p.probs))
 end
 
 """